from typing import Dict, Any, List
import aiohttp
from agents.components.embed.base_embedding_component import EmbeddingComponent
from agents.components.base import ComponentConfig

class OllamaEmbeddingComponent(EmbeddingComponent):
    """使用 Ollama 进行文本向量化的组件"""
    
    @classmethod
    def get_component_name(cls) -> str:
        return "ollama_embedding"
    
    @classmethod
    def get_component_description(cls) -> str:
        return "使用 Ollama 进行文本向量化的组件"
    
    @classmethod
    def get_config_schema(cls) -> Dict[str, Any]:
        base_schema = super().get_config_schema()
        base_schema["properties"].update({
            "ollama_url": {
                "type": "string",
                "description": "Ollama API 地址",
                "default": "http://localhost:11434"
            },
            "model_name": {
                "type": "string",
                "description": "Ollama 模型名称",
                "default": "nomic-embed-text"
            }
        })
        return base_schema
    
    def __init__(self, config: ComponentConfig):
        super().__init__(config)
        self.ollama_url = config.config.get("ollama_url", "http://localhost:11434")
        self.model_name = config.config.get("model_name", "nomic-embed-text")
        
    async def _embed_texts(self, texts: List[str]) -> Dict[str, Any]:
        embeddings = []
        
        async with aiohttp.ClientSession() as session:
            for text in texts:
                async with session.post(
                    f"{self.ollama_url}/api/embeddings",
                    json={
                        "model": self.model_name,
                        "prompt": text
                    }
                ) as response:
                    if response.status != 200:
                        raise Exception(f"Ollama API 调用失败: {await response.text()}")
                    
                    result = await response.json()
                    embeddings.append(result["embedding"])
        
        return {"embeddings": embeddings} 